import torchvision.transforms.functional as TF
from PIL import Image
import os

def erase_and_save(image_path, target_dir, position, size):
    """
       按照指定的位置和长宽擦除
       :param image_path: 输入图像路径
       :param target_dir: 目标图像路径
       :param position: 擦除的左上角坐标
       :param size: 擦除的长宽值
       :return: 返回擦除后的图像
    """
    image = TF.to_tensor(Image.open(image_path))
    erased_image = TF.to_pil_image(TF.erase(img=image,
                                            i=position[0],
                                            j=position[1],
                                            h=size[0],
                                            w=size[1],
                                            v=1))
    erased_image.save(os.path.join(target_dir, os.path.basename(image_path[:-4])) + "_erased_.jpg")
    return erased_image


def rotate_and_save(image_path, target_dir, angles_list):
    """
    按照一系列的指定角度旋转输入图像并保存
    :param image_path:输入图像路径
    :param target_dir:目标图片存储目录
    :param angles_list: 需要对输入图像进行多少角度的旋转
    :return: 完成后以列表形式返回所有的旋转后图像
    """
    image = Image.open(image_path)

    # 创建旋转后的图像容器
    image_list = []
    for angle in angles_list:
        rotated_image = TF.rotate(img=image, angle=angle, resample=Image.NEAREST)
        # 按照原始image_path的图片名称作为目标图像名称
        # 获取图片名称并去掉后缀".jpg"
        rotated_image.save(os.path.join(target_dir, os.path.basename(image_path[:-4])) + "_" + str(angle) + "_.jpg")
        image_list.append(rotated_image)
    return image_list


def vflip_and_save(image_path, target_path):
    """
    垂直旋转翻转图像，并保存
    :param image_path: 输入图像路径
    :param target_path: 目标图片存储目录，如"d/"
    :return: 返回旋转后的图片
    """
    image = Image.open(image_path)
    vertical_image = TF.vflip(img=image)
    vertical_image.save(os.path.join(target_path, os.path.basename(image_path[:-4]))+"_v_.jpg")

    return vertical_image

def gamma_and_save(image_path, target_dir, gamma_value):
    """
    进行伽马变换并保存
    :param image_path: 输入图片路径
    :param target_dir: 目标图像存储目录
    :param gamma_value: 伽马值
    :return: 伽马变换后的图像
    """
    image = Image.open(image_path)
    gamma_image = TF.adjust_gamma(img=image, gamma=gamma_value)
    gamma_image.save(os.path.join(target_dir, os.path.basename(image_path[:-4]))+"_gamma_.jpg")
    return gamma_image